package com.atguigu;

import org.apache.spark.SparkConf;
import org.apache.spark.sql.*;
import org.apache.spark.sql.expressions.Aggregator;

import java.io.Serializable;

import static org.apache.spark.sql.functions.udaf;
/**
 * @author yhm
 * @create 2022-12-26 9:54
 */
public class Test04_UDAF {
    public static void main(String[] args) {
        // 1. 创建sparkConf配置对象
        SparkConf conf = new SparkConf().setAppName("sql").setMaster("local[*]");

        // 2. 创建sparkSession连接对象
        SparkSession spark = SparkSession.builder().config(conf).getOrCreate();

        // 3. 编写代码
        Dataset<Row> dataset = spark.read().json("input/user.json");

        // 创建视图
        dataset.createOrReplaceTempView("t1");

        // 注册UDAF函数
        // 需要导入对应的udaf依赖
        spark.udf().register("myAvg",udaf(new MyAvg(),Encoders.LONG()));

        // 使用的效果
        spark.sql("select myAvg(age) avg_age from t1 ").show();

        // 4. 关闭sparkSession
        spark.close();
    }


    public static class Buffer implements Serializable {
        public Buffer() {
        }

        public Buffer(Long sum, Long count) {
            this.sum = sum;
            this.count = count;
        }

        private Long sum;
        private Long count;

        public Long getSum() {
            return sum;
        }

        public void setSum(Long sum) {
            this.sum = sum;
        }

        public Long getCount() {
            return count;
        }

        public void setCount(Long count) {
            this.count = count;
        }
    }

    // 父类的泛型
    // IN 传入的数据类型
    // BUF 计算的中间缓冲
    // OUT 输出的数据类型
    public static class MyAvg extends Aggregator<Long, Buffer, Double> {

        // 定义中间结果的初始值
        @Override
        public Buffer zero() {
            return new Buffer(0L, 0L);
        }

        // 中间结果和每一行的数值进行累加
        @Override
        public Buffer reduce(Buffer b, Long a) {
            b.setSum(b.getSum() + a);
            b.setCount(b.getCount() + 1);
            return b;
        }

        // 把不同分区的中间结果合并
        @Override
        public Buffer merge(Buffer b1, Buffer b2) {
            b1.setSum(b1.getSum() + b2.getSum());
            b1.setCount(b1.getCount() + b2.getCount());
            return b1;
        }

        // 根据聚合之后的中间结果计算最终值
        @Override
        public Double finish(Buffer reduction) {
            return reduction.getSum().doubleValue() / reduction.getCount();
        }

        @Override
        public Encoder<Buffer> bufferEncoder() {
            // 要求中间结果的类可以序列化  同时支持使用kryo序列化进行优化
            return Encoders.kryo(Buffer.class);
        }

        @Override
        public Encoder<Double> outputEncoder() {
            return Encoders.DOUBLE();
        }
    }
}
